import random
import numpy as np
import scipy.special as sc
from chaco.api import *
from traits.api import *
from traitsui.api import *
from chaco import default_colormaps
from enable.api import ComponentEditor
from chaco.tools.api import LineInspector, PanTool, ZoomTool

colormaps = list(default_colormaps.color_map_name_dict.keys())
for boring in "bone gray yarg gist_gray gist_yarg Greys".split():
    colormaps.remove(boring)


class Form(HasTraits):
    options = Enum("TM", "TE")
    min_x = CFloat(0)
    max_x = CFloat(360)
    npts_x = CInt(361)
    title_x = 'phi / deg'
    min_y = CFloat(.1)
    max_y = CFloat(20)
    npts_y = CInt(200)
    title_y = 'ka'
    value_label = 'RCS/dBm'
    xs = Array
    ys = Array
    zs = Dict
    plot_x = Instance(Plot)
    plot_y = Instance(Plot)
    container = Instance(HPlotContainer)
    polyplot = Instance(ContourPolyPlot)
    lineplot = Instance(ContourLinePlot)
    cross_plot_x = Instance(Plot)
    cross_plot_y = Instance(Plot)
    colorbar = Instance(ColorBar)
    num_levels = Int(15)
    colormap = Enum(colormaps)
    _image_index = Instance(GridDataSource)
    _image_value = Instance(ImageData)
    _cmap = Callable(default_colormaps.viridis)
    edit_model = Button
    xMaxIndex = Property(Int, observe=['xs'])
    yMaxIndex = Property(Int, observe=['ys'])
    xIndex = Int
    yIndex = Int

    def _get_xMaxIndex(self):
        return np.array(self.xs).size - 1

    def _get_yMaxIndex(self):
        return np.array(self.ys).size - 1

    def traits_init(self):
        self.compute_model()
        self.create_plot()
        self.create_plotxy()
        self.update()

    def compute_model(self):
        self.xs = np.linspace(self.min_x, self.max_x, self.npts_x)
        self.ys = np.linspace(self.min_y, self.max_y, self.npts_y)
        nRange = np.arange(-50, 51)
        ka, n = np.meshgrid(self.ys, nRange, indexing='ij')
        jn = sc.jv(n, ka)
        hn = sc.hankel2(n, ka)
        jn1 = sc.jv(n - 1, ka)
        hn1 = sc.hankel2(n - 1, ka)
        tm = (-1.) ** n * jn / hn / np.sqrt(np.pi * ka)
        te = (-1.) ** n * (ka * jn1 - n * jn) / (ka * hn1 - n * hn) / np.sqrt(np.pi * ka)
        n, phi = np.meshgrid(nRange, self.xs, indexing='ij')
        expValue = np.exp(1j * n * phi * np.pi / 180)
        tm = -2 * np.exp(1j * np.pi / 4) * tm @ expValue
        te = -2 * np.exp(1j * np.pi / 4) * te @ expValue
        self.zs['TM'] = 20 * np.log10(np.abs(tm))
        self.zs['TE'] = 20 * np.log10(np.abs(te))

    def create_plotxy(self):
        self.plot_x = Plot(ArrayPlotData())
        self.plot_y = Plot(ArrayPlotData())
        self.updatexy()
        self.plot_x.plot(('index', 'tm'), name='TM', color='auto')
        self.plot_x.plot(('index', 'te'), name='TE', color='auto')
        self.plot_x.x_axis.title = "Biscattering Angle / deg"
        self.plot_x.y_axis.title = "RCS/dBm"
        legend = Legend()
        legend.plots = self.plot_x.plots
        self.plot_x.overlays.append(legend)
        self.plot_y.plot(('index', 'tm'), name='TM', color='auto')
        self.plot_y.plot(('index', 'te'), name='TE', color='auto')
        self.plot_y.x_axis.title = "ka"
        self.plot_y.y_axis.title = "RCS/dBm"
        legend = Legend()
        legend.plots = self.plot_y.plots
        self.plot_y.overlays.append(legend)

    def create_plot(self):
        self._image_index = GridDataSource(np.array([]), np.array([]), sort_order=("ascending", "ascending"))
        image_index_range = DataRange2D(self._image_index)
        self._image_index.observe(self._metadata_changed, "metadata_changed")
        self._image_value = ImageData(data=np.array([]), value_depth=1)
        image_value_range = DataRange1D(self._image_value)
        self.polyplot = ContourPolyPlot(
            index=self._image_index,
            value=self._image_value,
            index_mapper=GridMapper(range=image_index_range),
            levels=self.num_levels,
            color_mapper=self._cmap(image_value_range),
        )
        self.lineplot = ContourLinePlot(
            index=self._image_index,
            value=self._image_value,
            index_mapper=GridMapper(range=self.polyplot.index_mapper.range),
            levels=self.num_levels,
        )
        self.left = PlotAxis(
            orientation="left",
            title=self.title_y,
            mapper=self.polyplot.index_mapper._ymapper,
            component=self.polyplot,
        )
        self.polyplot.overlays.append(self.left)
        self.bottom = PlotAxis(
            orientation="bottom",
            title=self.title_x,
            mapper=self.polyplot.index_mapper._xmapper,
            component=self.polyplot,
        )
        self.polyplot.overlays.append(self.bottom)
        self.polyplot.tools.append(
            PanTool(self.polyplot, constrain_key="shift")
        )
        self.polyplot.overlays.append(
            ZoomTool(component=self.polyplot, tool_mode="box", always_on=False)
        )
        self.polyplot.overlays.append(
            LineInspector(
                component=self.polyplot,
                axis="index_x",
                inspect_mode="indexed",
                write_metadata=True,
                color="white",
                is_listener=True,
            )
        )
        self.polyplot.overlays.append(
            LineInspector(
                component=self.polyplot,
                axis="index_y",
                inspect_mode="indexed",
                write_metadata=True,
                color="white",
                is_listener=True,
            )
        )
        contour_container = OverlayPlotContainer(padding=20, use_backbuffer=True, unified_draw=True)
        contour_container.add(self.polyplot)
        contour_container.add(self.lineplot)
        cbar_index_mapper = LinearMapper(range=image_value_range)
        self.colorbar = ColorBar(
            index_mapper=cbar_index_mapper,
            plot=self.polyplot,
            padding_top=self.polyplot.padding_top,
            padding_bottom=self.polyplot.padding_bottom,
            padding_right=40,
            resizable="v",
            width=30,
        )
        self.cross_plot_x = Plot(
            ArrayPlotData(
                line_index=np.array([]),
                line_value=np.array([]),
                scatter_index=np.array([]),
                scatter_value=np.array([]),
                scatter_color=np.array([]),
            ),
            resizable="h",
            height=140,
            padding=20,
        )
        self.cross_plot_x.plot(("line_index", "line_value"), line_style="dot")
        self.cross_plot_x.plot(
            ("scatter_index", "scatter_value", "scatter_color"),
            type="cmap_scatter",
            name="dot",
            color_mapper=self._cmap(image_value_range),
            marker="circle",
            marker_size=8,
        )
        self.cross_plot_x.index_range = self.polyplot.index_range.x_range
        self.cross_plot_y = Plot(
            ArrayPlotData(
                line_index=np.array([]),
                line_value=np.array([]),
                scatter_index=np.array([]),
                scatter_value=np.array([]),
                scatter_color=np.array([]),
            ),
            width=140,
            resizable="v",
            padding=20,
            orientation="v",
            padding_bottom=200,
        )
        self.cross_plot_y.plot(("line_index", "line_value"), line_style="dot")
        self.cross_plot_y.plot(
            ("scatter_index", "scatter_value", "scatter_color"),
            type="cmap_scatter",
            name="dot",
            color_mapper=self._cmap(image_value_range),
            marker="circle",
            marker_size=8,
        )
        self.cross_plot_y.index_range = self.polyplot.index_range.y_range
        self.container = HPlotContainer(
            padding=40,
            fill_padding=True,
            bgcolor="white",
            use_backbuffer=False,
        )
        inner_cont = VPlotContainer(padding=0, use_backbuffer=True)
        inner_cont.add(self.cross_plot_x)
        inner_cont.add(contour_container)
        self.container.add(self.colorbar)
        self.container.add(inner_cont)
        self.container.add(self.cross_plot_y)

    @observe('options')
    def update(self, event=None):
        self.colorbar.index_mapper.range.low = np.nanmin(self.zs[self.options])
        self.colorbar.index_mapper.range.high = np.nanmax(self.zs[self.options])
        self._image_index.set_data(self.xs, self.ys)
        self._image_value.data = self.zs[self.options]
        self.cross_plot_x.data.update_data(line_index_x=self.xs)
        self.cross_plot_y.data.update_data(line_index_y=self.ys)

    def _metadata_changed(self, event=None):
        if "selections" in self._image_index.metadata and len(self._image_index.metadata['selections']) == 2:
            x_ndx, y_ndx = self._image_index.metadata["selections"]
            if y_ndx and x_ndx:
                xdata, ydata = self._image_index.get_data()
                xdata, ydata = xdata.get_data(), ydata.get_data()
                self.cross_plot_x.data.update_data(
                    line_index=self.xs,
                    line_value=self._image_value.data[y_ndx, :],
                    scatter_index=np.array([xdata[x_ndx]]),
                    scatter_value=np.array([self._image_value.data[y_ndx, x_ndx]]),
                    scatter_color=np.array([self._image_value.data[y_ndx, x_ndx]]),
                )
                self.cross_plot_y.data.update_data(
                    line_index=self.ys,
                    line_value=self._image_value.data[:, x_ndx],
                    scatter_index=np.array([ydata[y_ndx]]),
                    scatter_value=np.array([self._image_value.data[y_ndx, x_ndx]]),
                    scatter_color=np.array([self._image_value.data[y_ndx, x_ndx]]),
                )
        else:
            self.cross_plot_x.data.update_data(
                line_value=np.array([]),
                scatter_value=np.array([]),
            )
            self.cross_plot_y.data.update_data(
                line_value=np.array([]),
                scatter_value=np.array([]),
            )

    def _colormap_changed(self):
        self._cmap = default_colormaps.color_map_name_dict[self.colormap]
        if self.polyplot is not None:
            value_range = self.polyplot.color_mapper.range
            self.polyplot.color_mapper = self._cmap(value_range)
            value_range = self.cross_plot_x.color_mapper.range
            self.cross_plot_x.color_mapper = self._cmap(value_range)
            self.cross_plot_y.color_mapper = self._cmap(value_range)
        self.container.request_redraw()

    def _num_levels_changed(self):
        if self.num_levels > 3:
            if self.polyplot is not None:
                self.polyplot.levels = self.num_levels
            if self.lineplot is not None:
                self.lineplot.levels = self.num_levels

    @observe('xIndex,yIndex')
    def updatexy(self, event=None):
        self.plot_x.data.update_data(
            index=self.xs,
            tm=self.zs['TM'][self.yIndex, :],
            te=self.zs['TE'][self.yIndex, :],
        )
        self.plot_y.data.update_data(
            index=self.ys,
            tm=self.zs['TM'][:, self.xIndex],
            te=self.zs['TE'][:, self.xIndex],
        )
        self.plot_x.title = f'ka={self.ys[self.yIndex]}'
        self.plot_y.title = f'Bistatic Angle={self.xs[self.xIndex]} deg'

    class ModelHandler(Handler):
        def closed(self, info, is_ok):
            if is_ok:
                obj = info.object
                obj.compute_model()
                obj.update()
                obj.updatexy()
            else:
                print('Cancel')

    model_view = View(
        Group(
            HGroup(
                Item("npts_x", label="双站角点数"),
                Item("npts_y", label="ka点数"),
            ),
            HGroup(
                Item("min_x", label="最小双站角"),
                Item("min_y", label="最小ka"),
            ),
            HGroup(
                Item("max_x", label="最大双站角"),
                Item("max_y", label="最大ka"),
            ),
        ),
        buttons=["OK", "Cancel"],
        handler=ModelHandler,
    )
    traits_view = View(
        Group(
            Tabbed(
                Group(
                    UItem("container", editor=ComponentEditor(size=(800, 600))),
                    HGroup(
                        Item("num_levels", label='等高线数量'),
                        Item("colormap", label='颜色'),
                        Item("@options", label='选项'),
                        spring,
                    ),
                    label='ka、双站角'
                ),
                Group(
                    Group(UItem("plot_x", editor=ComponentEditor(size=(800, 600)))),
                    Item('yIndex', editor=RangeEditor(
                        low=0,
                        high_name='yMaxIndex',
                        mode='slider')),
                    label='双站角'
                ),
                Group(
                    Group(UItem("plot_y", editor=ComponentEditor(size=(800, 600)))),
                    Item('xIndex', editor=RangeEditor(
                        low=0,
                        high_name='xMaxIndex',
                        mode='slider')),
                    label='ka'
                ),
            ),
            HGroup(
                UItem("edit_model", label='设置参数'),
                spring,
            )
        ),
        title="导体柱散射",
        resizable=True,
    )

    def _edit_model_fired(self):
        self.configure_traits(view='model_view')


if __name__ == "__main__":
    form = Form()
    form.configure_traits()
